#!/usr/bin/env python
# coding: utf-8

# # 训练和评估
# 
# 前面章节讲解了luojianet_ms构建网络所使用的基本元素，如luojianet_ms的网络基本单元、损失函数、优化器和评价函数等。
# 
# 本章重点介绍如何使用这些元素自定义训练和评估网络。

# ## 构建训练和评估
# 
# 构建训练网络首先需要构建前向网络，然后在前向网络的基础上叠加损失函数、反向传播和优化器。
# 
# ### 定义数据集
# 
# 如下示例定义`get_data`函数生成样本数据及对应的标签，定义`create_dataset`函数加载自定义数据集。

# In[5]:


import luojianet_ms.dataset as ds
import numpy as np

def get_data(num, w=2.0, b=3.0):
    """生成样本数据及对应的标签"""
    for _ in range(num):
        x = np.random.uniform(-10.0, 10.0)
        noise = np.random.normal(0, 1)
        y = x * w + b + noise
        yield np.array([x]).astype(np.float32), np.array([y]).astype(np.float32)

def create_dataset(num_data, batch_size=16):
    """生成数据集"""
    dataset = ds.GeneratorDataset(list(get_data(num_data)), column_names=['data', 'label'])
    dataset = dataset.batch(batch_size)
    return dataset


# ### 构建前向网络
# 
# 使用`nn.Module`构建前向网络，如下示例定义一个简单的线性回归网络`LinearNet`：

# In[6]:


import numpy as np
import luojianet_ms.nn as nn
from luojianet_ms.common.initializer import Normal

class LinearNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc = nn.Dense(1, 1, Normal(0.02), Normal(0.02))

    def forward(self, x):
        return self.fc(x)


# ### 构建训练流程
# 
# luojianet_ms的`nn`模块提供了训练网络封装函数`TrainOneStepCell`，用来封装网络和优化器。其参数如下：
# 
# - `network`：训练网络，只支持单输出网络。
# - `optimizer`： 用于更新网络参数的优化器。
# - `sens`：反向传播的输入，缩放系数，默认值为1.0。
# 
# 如下示例使用`nn.TrainOneStepCell`将上述定义的线性回归网络封装成一个训练网络，并执行训练，打印损失值。
# 
# 示例代码中使用`set_train`通过`mode`参数指定模型是否为训练模式，其中`mode`参数默认为True，即默认情况下为训练模式，若`mode`为False，则为评估或推理模式。

# In[7]:


# 生成训练数据集
train_dataset = create_dataset(num_data=160, batch_size=16)

net = LinearNet()
loss = nn.MSELoss()

# 连接前向网络与损失函数
net_with_loss = nn.WithLossCell(net, loss)
opt = nn.Momentum(net.trainable_params(), learning_rate=0.005, momentum=0.9)

# 定义训练网络，封装网络和优化器
train_net = nn.TrainOneStepCell(net_with_loss, opt)
# 设置网络为训练模式
train_net.set_train()

# 真正训练迭代过程
step = 0
epochs = 2
steps = train_dataset.get_dataset_size()

for epoch in range(epochs):
    for d in train_dataset.create_dict_iterator():
        result = train_net(d["data"], d["label"])
        print(f"Epoch: [{epoch} / {epochs}], "
              f"step: [{step} / {steps}], "
              f"loss: {result}")
        step = step + 1


# ### 构建评估流程
# 
# luojianet_ms的`nn`模块提供了评估网络封装函数`WithEvalCell`，用来在验证集上评估模型训练的效果。其参数如下：
# 
# - `network`：前向网络。
# - `loss_fn`：损失函数。
# - `add_cast_fp32`：是否将数据类型调整为float32。
# 
# `nn.WithEvalCell`只接受两个输入，分别为数据`data`及其对应的标签`label`，用前面定义的前向网络和损失函数构建一个评估网络，示例如下：

# In[8]:


eval_dataset = create_dataset(num_data=160, batch_size=16)

# 构建评估网络
eval_net = nn.WithEvalCell(net, loss)
eval_net.set_train(False)
loss = nn.Loss()
mae = nn.MAE()

mae.clear()
loss.clear()

# 真正验证迭代过程
for data in eval_dataset.create_dict_iterator():
    outputs = eval_net(data["data"], data["label"])
    mae.update(outputs[1], outputs[2])
    loss.update(outputs[0])

# 评估结果
mae_result = mae.eval()
loss_result = loss.eval()

print("mae: ", mae_result)
print("loss: ", loss_result)


# ## 自定义训练和评估
# 
# ### 自定义训练网络
# 
# 自定义损失函数章节已经介绍了使用`nn.WithLossCell`将前向网络与损失函数连接起来，本节将介绍如何自定义训练网络。
# 
# 如下示例定义`CustomTrainOneStepCell`函数来封装网络和优化器。

# In[10]:


import luojianet_ms.ops as ops

class CustomTrainOneStepCell(nn.Module):
    """自定义训练网络"""

    def __init__(self, network, optimizer):
        """入参有两个：训练网络，优化器"""
        super(CustomTrainOneStepCell, self).__init__(auto_prefix=False)
        self.network = network                           # 定义前向网络
        self.network.set_grad()                          # 构建反向网络
        self.optimizer = optimizer                       # 定义优化器
        self.weights = self.optimizer.parameters         # 待更新参数
        self.grad = ops.GradOperation(get_by_list=True)  # 反向传播获取梯度

    def forward(self, *inputs):
        loss = self.network(*inputs)                            # 计算当前输入的损失函数值
        grads = self.grad(self.network, self.weights)(*inputs)  # 进行反向传播，计算梯度
        self.optimizer(grads)                                   # 使用优化器更新权重参数
        return loss

net1 = LinearNet()   # 定义前向网络
loss = nn.MSELoss()  # 损失函数


# 连接前向网络与损失函数
net_with_loss = nn.WithLossCell(net1, loss)
opt = nn.Momentum(net1.trainable_params(), learning_rate=0.005, momentum=0.9)

# 定义训练网络，封装网络和优化器
train_net = CustomTrainOneStepCell(net_with_loss, opt)
# 设置网络为训练模式
train_net.set_train()

# 真正训练迭代过程
step = 0
epochs = 2
steps = train_dataset.get_dataset_size()
for epoch in range(epochs):
    for d in train_dataset.create_dict_iterator():
        result = train_net(d["data"], d["label"])
        print(f"Epoch: [{epoch} / {epochs}], "
              f"step: [{step} / {steps}], "
              f"loss: {result}")
        step = step + 1


# ### 自定义评估网络
# 
# 由于`nn.WithEvalCell`只有两个输入`data`和`label`，对于多个数据或多个标签的场景显然不适用，此时如果想要构建评估网络就需要自定义评估网络。
# 
# 在自定义时，如不需要损失函数作为评价指标，则无需定义`loss_fn`。当输入为多数据或多标签时，可参考如下示例来自定义评估网络。

# In[11]:


class CustomWithEvalCell(nn.Module):

    def __init__(self, network):
        super(CustomWithEvalCell, self).__init__(auto_prefix=False)
        self.network = network

    def forward(self, data, label1, label2):
        """输入数据为三个：一个数据及其对应的两个标签"""
        outputs = self.network(data)
        return outputs, label1, label2

custom_eval_net = CustomWithEvalCell(net)
custom_eval_net.set_train(False)


# ## 网络的权重共享
# 
# 通过前面的介绍可以看出，前向网络、训练网络和评估网络具有不同的逻辑，因此在需要时我们会构建三张网络。若使用训练好的模型进行评估和推理，需要推理和评估网络中的权重值与训练网络中相同。
# 
# 使用模型保存和加载接口，将训练好的模型保存下来，再加载到评估和推理网络中，可以确保权重值相同。在训练平台上完成模型训练，但在推理平台进行推理时，模型保存与加载是必不可少的。
# 
# 在网络调测过程中，或使用边训练边推理方式进行模型调优时，往往在同一Python脚本中完成模型训练，评估或推理，此时luojianet_ms的权重共享机制可确保不同网络间的权重一致性。
# 
# 使用luojianet_ms构建不同网络结构时，只要这些网络结构是在一个实例的基础上封装的，那这个实例中的所有权重便是共享的，一个网络中的权重发生变化，意味着其他网络中的权重同步发生了变化。
# 
# 如下示例中，定义训练和评估网络时便使用了权重共享机制：

# In[12]:


# 实例化前向网络
net = LinearNet()
# 设定损失函数并连接前向网络与损失函数
loss = nn.MSELoss()
net_with_loss = nn.WithLossCell(net, loss)
# 设定优化器
opt = nn.Adam(params=net.trainable_params())

# 定义训练网络
train_net = nn.TrainOneStepCell(net_with_loss, opt)
train_net.set_train()

# 构建评估网络
eval_net = nn.WithEvalCell(net, loss)
eval_net.set_train(False)


# `train_net`和`eval_net`均在`net`实例的基础上封装，因此在进行模型评估时，不需要加载`train_net`的权重。
# 
# 若在构建`eval_net`时重新的定义前向网络，那`train_net`和`eval_net`之间便没有共享权重，如下：

# In[13]:


# 定义训练网络
train_net = nn.TrainOneStepCell(net_with_loss, opt)
train_net.set_train()

# 再次实例化前向网络
net2 = LinearNet()
# 构建评估网络
eval_net = nn.WithEvalCell(net2, loss)
eval_net.set_train(False)


# 此时，若要在模型训练后进行评估，就需要将`train_net`中的权重加载到`eval_net`中。在同一脚本中进行模型训练、评估和推理时，利用好权重共享机制不失为一种更简便的方式。
# 
# 在使用Model进行训练时，对于简单的场景，`Model`内部使用`nn.WithLossCell`、`nn.TrainOneStepCell`和`nn.WithEvalCell`在前向`network`实例的基础上构建训练和评估网络，`Model`本身确保了推理、训练、评估网络之间权重共享。
# 
# 但对于自定义使用Model的场景，用户需要注意前向网络仅实例化一次。如果构建训练网络和评估网络时分别实例化前向网络，那在使用`eval`进行模型评估时，便需要手动加载训练网络中的权重，否则模型评估使用的将是初始的权重值。
